package io.modelcontextprotocol.autoconfigure.definition;


import com.fasterxml.jackson.annotation.JsonProperty;
import com.github.victools.jsonschema.generator.FieldScope;
import com.github.victools.jsonschema.generator.MemberScope;
import com.github.victools.jsonschema.generator.Module;
import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder;
import com.github.victools.jsonschema.generator.SchemaGeneratorConfigPart;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.media.Schema.RequiredMode;
import java.util.stream.Stream;
import org.springframework.lang.Nullable;
import org.springframework.util.StringUtils;

public final class SpringAiSchemaModule implements Module {
    private final boolean requiredByDefault;

    public SpringAiSchemaModule(Option... options) {
        this.requiredByDefault = Stream.of(options).noneMatch((option) -> {
            return option == SpringAiSchemaModule.Option.PROPERTY_REQUIRED_FALSE_BY_DEFAULT;
        });
    }

    public void applyToConfigBuilder(SchemaGeneratorConfigBuilder builder) {
        this.applyToConfigBuilder(builder.forFields());
    }

    private void applyToConfigBuilder(SchemaGeneratorConfigPart<FieldScope> configPart) {
        configPart.withDescriptionResolver(this::resolveDescription);
        configPart.withRequiredCheck(this::checkRequired);
    }

    @Nullable
    private String resolveDescription(MemberScope<?, ?> member) {
        ToolParam toolParamAnnotation = (ToolParam)member.getAnnotationConsideringFieldAndGetter(ToolParam.class);
        return toolParamAnnotation != null && StringUtils.hasText(toolParamAnnotation.description()) ? toolParamAnnotation.description() : null;
    }

    private boolean checkRequired(MemberScope<?, ?> member) {
        ToolParam toolParamAnnotation = (ToolParam)member.getAnnotationConsideringFieldAndGetter(ToolParam.class);
        if (toolParamAnnotation != null) {
            return toolParamAnnotation.required();
        } else {
            JsonProperty propertyAnnotation = (JsonProperty)member.getAnnotationConsideringFieldAndGetter(JsonProperty.class);
            if (propertyAnnotation != null) {
                return propertyAnnotation.required();
            } else {
                Schema schemaAnnotation = (Schema)member.getAnnotationConsideringFieldAndGetter(Schema.class);
                if (schemaAnnotation == null) {
                    Nullable nullableAnnotation = (Nullable)member.getAnnotationConsideringFieldAndGetter(Nullable.class);
                    return nullableAnnotation != null ? false : this.requiredByDefault;
                } else {
                    return schemaAnnotation.requiredMode() == RequiredMode.REQUIRED || schemaAnnotation.requiredMode() == RequiredMode.AUTO || schemaAnnotation.required();
                }
            }
        }
    }

    public static enum Option {
        PROPERTY_REQUIRED_FALSE_BY_DEFAULT;

        private Option() {
        }
    }
}
